import pandas as pd
import numpy as np


def extract_dt(df):
    df['time'] = pd.to_datetime(df['time'], format='%m%d %H:%M:%S')
    df['date'] = df['time'].dt.date
    df['hour'] = df['time'].dt.hour
    df['weekday'] = df['time'].dt.weekday
    return df


def diff_op(df, cols_sp):
    def diff(a, cols_sp):
        col = a.name
        if col not in cols_sp:
            arr = a.tolist()
            b = pd.Series(arr[1:] + [arr[-1]]) - a
        else:
            b = a
        return b
    df1 = df.sort_values(by='time').reset_index().drop('index', axis=1)
    df2 = df1.apply(diff, axis=0, **{'cols_sp': cols_sp})
    return df2


train = pd.read_hdf('./input/train.h5')
test = pd.read_hdf('./input/test.h5')
train = extract_dt(train)
test = extract_dt(test)

ship_type = train.drop_duplicates(['ship', 'type'])[['ship', 'type']]
type_map = {'围网': 0, '拖网': 1, '刺网': 2}
type_map_rev = {v: k for k, v in type_map.items()}
ship_type['type_label'] = ship_type['type'].map(type_map)


train['rnk'] = range(train.shape[0])
cols = ['x', 'y', 'v', 'd', 'time', 'rnk', 'ship']
feas = ['x', 'y', 'v', 'd', 'time', 'rnk']
cols_sp = ['rnk']
trn_ = train[cols].groupby('ship')[feas].apply(diff_op, cols_sp).\
    reset_index().drop('level_1', axis=1)
trn_.columns = ['ship_1', 'xd', 'yd', 'vd', 'dd', 'td', 'rnk']

trn1 = pd.merge(train, trn_, on='rnk', suffixes=['', ''])
trn1['sd'] = trn1['td'].dt.seconds
trn1['date'] = trn1['time'].dt.day
trn1['hour'] = trn1['time'].dt.hour
trn1['month'] = trn1['time'].dt.month

drop_cols = ['rnk', 'time', 'rnk', 'ship_1']
trn2 = trn1.drop(drop_cols, axis=1)


trn = trn2.copy()
trn['xd_x_yd'] = abs(trn['xd'])*abs(trn['yd'])
trn['xd_yd'] = pd.Series((trn['xd']**2 + trn['yd']**2)**0.5)
trn['yd/xd'] = trn['yd']/trn['xd']
trn['xd/yd'] = trn['xd']/trn['yd']
trn['xd/s'] = trn['xd']/trn['sd']
trn['yd/s'] = trn['yd']/trn['sd']
trn['vd/s'] = trn['vd']/trn['sd']
trn['dd/s'] = trn['dd']/trn['sd']
trn['xd_x_yd/s'] = trn['xd_x_yd']/trn['sd']
trn['xd_yd/s'] = trn['xd_yd']/trn['sd']
trn['v_x_s'] = trn['v']*trn['sd']
trn['d_x_s'] = trn['d']*trn['sd']
trn['xd_x_s'] = (trn['xd'] != 0)*trn['sd']
trn['yd_x_s'] = (trn['yd'] != 0)*trn['sd']
trn.replace({np.nan: 0, np.inf: 0, -np.inf: 0}, inplace=True)

cols = ['date', 'hour', 'weekday', 'xd',
       'yd', 'vd', 'dd', 'sd', 'month', 'xd_x_yd', 'xd_yd', 'yd/xd',
       'xd/yd', 'xd/s', 'yd/s', 'vd/s', 'dd/s', 'xd_x_yd/s', 'xd_yd/s',
       'v_x_s', 'd_x_s', 'xd_x_s', 'yd_x_s', 'ship', 'type']
trn[cols].to_hdf('input/train_d.h5', 'df', mode='w')




# 所以0的个数也要计算，和也要计算
# (trn[cols]==0).sum(axis=0)/trn.shape[0]*100
# ship        0.015335
# x           0.000000
# y           0.000000
# v          20.307241
# d          32.154756
# time        0.000000
# type        0.000000
# date        0.000000
# hour        4.197341
# weekday    17.292837
# rnk         0.000037
# ship_1      0.015335
# xd         51.026841
# yd         51.026767
# vd         25.286020
# dd         29.267702
# td          0.000000
# sd          0.259294
# month       0.000000
# trn_.loc[trn_['ship_1']==0, :]
# train.loc[trn_['ship']==0, :]